package com.ai.rag.qa.springaialibabaragqa.assistant;


import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.chat.prompt.SystemPromptTemplate;
import org.springframework.ai.document.Document;
import org.springframework.stereotype.Component;

import java.util.List;
import java.util.Map;

@Component
public class AiAssistant {

    private final ChatClient dashScopeChatClient;

    private final String DEFAULT_SYSTEM_PROMPT = "你是一个答疑机器人，你的任务是根据下述给定的已知信息回答用户的问题 已知信息: {context}\n" +
//            "    用户问题: {question}\n" +
            "    如果已知信息不包含用户问题的答案，或者已知信息不足以回答用户的问题，请直接回复“我无法回答您的问题”.\n" +
            "    请不要输出已知信息中不包含的信息或者答案。\n" +
            "    请用中文回答用户问题";

    public AiAssistant(ChatModel chatModel) {
        this.dashScopeChatClient = ChatClient.builder(chatModel).defaultSystem(DEFAULT_SYSTEM_PROMPT).build();
    }

    public String chat(String message, List<Document> context){
        SystemPromptTemplate systemPromptTemplate = new SystemPromptTemplate(DEFAULT_SYSTEM_PROMPT);
        Message systemMessage = systemPromptTemplate.createMessage(Map.of("context", context));
        UserMessage userMessage = new UserMessage( message);

        return this.dashScopeChatClient
                .prompt(new Prompt(List.of(systemMessage, userMessage)))
                .call().content();
    }
}